#ifndef TERN_SCHED_CACHE_H
#define TERN_SCHED_CACHE_H

#include <sys/syscall.h> 
#include <sys/types.h>     
#include <errno.h>
#include <unistd.h>
#include <sys/time.h>
#include <pthread.h>
#include <sys/errno.h>          
#include <time.h>
#include <sys/time.h> 
#include <iostream>
#include <fstream>
#include <sstream>
#include <ext/hash_map>
#include <set>

#include "expr/Lexer.h"
#include "expr/Parser.h"

#include "klee/Constraints.h"
#include "klee/Expr.h"
#include "klee/ExprBuilder.h"
#include "klee/Solver.h"
#include "klee/Statistics.h"
#include "klee/util/ExprPPrinter.h"
#include "klee/util/ExprVisitor.h"
#include "klee/util/ExprUtil.h"
#include "klee/util/Assignment.h"
#include "klee/Internal/ADT/MapOfSets.h"
#include "tern/TernDefs.h"

#include "llvm/ADT/OwningPtr.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/ManagedStatic.h"
#include "llvm/Support/MemoryBuffer.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/System/Signals.h"
#include "llvm/Function.h"

#define sc_print(args...) //{fprintf(stderr, args); fflush(stderr);}
#define my_print(args...) //{fprintf(stderr, args); fflush(stderr);}

using namespace llvm;
using namespace std;
using namespace klee::expr;

/*	Below are all internel structures or ulitity functions (e.g., merging constraints with "OR" for the same schedule) for the schedule cache.
	Iterator to the schedule cache? This is important to the replaying part. This is also internal. */

namespace klee {
	class SchedTrunk;
	class SchedNode;
	class MemoizedSched;
	class CacheItr;
	//class MemoizedSchedule;
	class SchedCache;

	typedef std::set<ref<Expr> > KeyType;

	// An element of a schedule.
	class SchedTrunk {
	private:
		// The first trunk Id denotes the "track".
		unsigned trunkSize;
		int trunkIds[SCHED_TRUNK_MAX_SIZE];
		char *trunkInfo[SCHED_TRUNK_MAX_SIZE];
		int ops[SCHED_TRUNK_MAX_SIZE];
		//vector<int> trunkIds;														// A trunk of schedule, includes only tids.
		//vector<const char *> trunkInfo;										// A trunk of schedule, includes all information, including the landmark name for each id.
		ConstraintManager *cManager;									// Store constraints for this particular symbolic var along this trunk, so constraints are binded with sched trunk.
		std::vector<const Array*> objects;
		SchedNode *prevNode;
		SchedNode *nextNode;
		//int instId;																		// Instruction id of the "tern_symbolic" llvm instruction call.
		std::string tidFileName;
		std::string cstFileName;
		Char_Ptr_To_Int_Map cache;
		std::vector<unsigned> enabledTracks;

	protected:
	
	public:
		SchedTrunk();
		~SchedTrunk();
		std::string& getTidFileName();
		std::string& getCstFileName();
		std::string& dumpTidsToDisk(unsigned index, int winNum);
		void loadTidsFromDisk(std::string name);
		std::string& dumpConstraintsToDisk(unsigned index, int winNum);
		void loadConstraintsFromDisk(std::string name);
		ConstraintManager *parseConstraints(const char *Filename, const MemoryBuffer *MB, ExprBuilder *Builder);
		unsigned size();
		SchedNode *getPrevNode();
		void setPrevNode(SchedNode *node);
		SchedNode *getNextNode();
		void setNextNode(SchedNode *node);
		int tidAt(unsigned index);
		int setTidAt(unsigned index, int newTid);
		char *tidInfoAt(unsigned index);
		void setTidInfoAt(unsigned index, char *newTidInfo);
		int opAt(unsigned index);
		ConstraintManager *getCst();
		std::vector<const Array*> &getObjects();
		Char_Ptr_To_Int_Map &getCache();
		unsigned constraintSize(unsigned trackId);
		int getOpByTidInfo(char *tidInfo);
		void updateTidInfoWithNewTid(unsigned index, int newTid);
		//int getTidFromTidInfo(char *tidInfo);
		int addTid(char *tidInfo);
		int addConstraint(unsigned trackId, ref<Expr> expr);
		void addSymVar(const Array *arr);
		bool isLeafTrunk();
		void addEnabledTrack(unsigned trackIdx);
		std::vector<unsigned> &getEnabledTracks();
	};

	// This class can be used in both recording mode and replaying mode.
	// In recording mode, this class starts and ends with taskStart and taskEnd.
	// At each taskEnd, a schedule is merged into the decision tree (the decision tree is used in only mode), and the schedule is not destroyed.
	// Schedules are still stored linearly.
	// In replaying mode, this class passes each linear memoized schedule (just like each schedule is memoized) 
	// and rebuild the decision tree.
	class MemoizedSched {
	private:
		vector<SchedTrunk *> schedTrunks;

	public:
		MemoizedSched();
		void dumpSchedTrunksToDisk(int winNum);
		unsigned size();
		SchedTrunk *getTrunk(unsigned index);
		int addTid(char *tidInfo);
		int addConstraint(unsigned trackId, ref<Expr> expr);
		void addSymVar(const Array *arr);
		void normalize();
	};
	
	// Each schedule node corresponds to a runtime instance of "tern_symbolic", not statically.
	// This class only involves with SchedCache, not MemoizedSchedule.
	class SchedNode {
	private:
		SchedTrunk *prevTrunk;
		
		// WARN: when building schedule cache decision tree, need to make sure each child trunk has a set of "distinct" constraints.
		// I.e., no two or more set of constraints in different trunks can be satisfied with the same concrete value.
		// If constraints from a new schedule is not distinct with existing constraints, then just reuse the existing trunks and ignore the new trunks.
		// Otherwise, for a new set of concrete input, there could be multiple trunks at a sched node that can satisfy this input.
		vector<SchedTrunk *> children;
	public:
		SchedNode();
		// We will merge two sets of constraints to one SchedNode iff their constraints and sched trunks are distinct.
		void addChild(SchedTrunk *schedTrunk);
		SchedTrunk *getPrevTrunk();
		void setPrevTrunk(SchedTrunk *st);
		SchedTrunk *getChild(unsigned index);
		void delChild(unsigned index);
		unsigned size();
	};

	// The Sched Cache Iterator to control replay. This should not be a field of the SchedCache since it contains run time thread information.
	// It should be part of the replayer.
	class CacheItr {
	private:		
		/* Flags that can be maintained by atomic operations. */
		SchedNode *curNode;
		SchedTrunk *curTrunk;
		volatile int offset;
		volatile unsigned move;

		/* Race-free arrays. */
		int tidMap[MAX_THREAD_NUM];						// Map from replayer maintained thread id to the schedule template thread id.
		int reversedTidMap[MAX_THREAD_NUM];
		struct timeval tInit;
		struct timeval tNow;

		// For sched track mode.
		int curSchedTrackIdx;
		unsigned trackTrunkOffset;

	protected:

	public:
		CacheItr();
		int getOffset();
		unsigned getMove();
		SchedTrunk *getCurTrunk();
		SchedNode *getCurNode();
		void setCurNode(SchedNode *sn);
		void setCurTrunk(SchedTrunk *st);
		bool myTurn(unsigned arg, int templateTid, int hookId);
		MoveOnRet moveOn(int tid, int hookId, unsigned arg, int &nextTid);
		void setTidMap(int replayerTid, int templateTid);
		int getTidMap(int replayerTid);
		int getReversedTidMap(int templateTid);
		void init();
		void moveItrToNextTrunk(int replayTid, SchedTrunk *nextTrunk);
		char *fillSymVarMap(const char *name, void *addr, unsigned nBytes);
		void initSchedTrack(unsigned idx, SchedTrunk *cT);
		unsigned moveOnSchedTrackTrunk();
		unsigned getSchedTrackIdx() {return curSchedTrackIdx;};
		Sym_Var_To_Value_Map symVarValueMap;
		volatile int hTplTid;
	};

	class SchedTrack {
	private:
		const char *schedId;								// The last trunk of a schedule, i.e., the leaf of a path to identify a schedule.
		const char *keyWord;								// to select expressions from a schedule, like "sqlcmd".
		std::vector< std::vector<ConstraintManager *> > trackManager;
		std::vector< std::vector<const Array*> > trackObjects;
		std::vector< std::vector<int> > replayTidWaitQueue;
		std::vector<SchedTrunk *> trunks;
		unsigned index;

	public:
		SchedTrack(const char *id, const char *kw, unsigned index);
		void addConstraint(unsigned trackId, ref <Expr> expr);
		void buildTrack(SchedTrunk *cT, SchedNode *root);
		bool checkTrackAt(unsigned trackId, std::vector< std::vector<unsigned char> > &values, int checkerReplayTid);
		std::vector<ConstraintManager *> &getCst(unsigned trackId);
		bool check(ConstraintManager *cm, std::vector<const Array*> &objects, std::vector< std::vector<unsigned char> > &values);
		void addCManager(unsigned trackId, ConstraintManager *cm);
		void addObj(unsigned trackId, const Array *arr);
		std::vector<ConstraintManager *> &getCsts(unsigned trackId);
		std::vector<const Array *> &getObjs(unsigned trackId);
		unsigned getBackCManagerSize(unsigned trackId);
		const char *getTrackId();
		void insertReplayTidAt(unsigned trackId, int replayTid);
		bool isFull();
		int start(unsigned trackId);
		void delTid(int replayTid);
		bool isEmptyAt(unsigned trackId);
		unsigned emptyNum();
		unsigned getIndex();
		SchedTrunk *getTrunkAt(unsigned index);
	};

	// The cache built during the recording time. After a schedule is being recoded, it will be merged into the schedule cache.
	// This is a decision tree.
	class SchedCache {
	private:
		Solver *solver;
		unsigned schedNum;
		typedef hash_map<int, MemoizedSched *> Int_To_MemoizedSched_Map;
		Int_To_MemoizedSched_Map schedules;
		SchedNode *root;
		SchedNode *itr;

		std::vector <SchedTrack *> tracks;
		unsigned numTracks;

		// Old implementation.
		//typedef hash_map<int, MemoizedSchedule *> Int_To_MemoizedSchedule_Map;
		//Int_To_MemoizedSchedule_Map constraintSet;

	protected:
		Solver *constructSolverChain(STPSolver *stpSolver, std::string queryLogPath, 
		std::string stpQueryLogPath, std::string queryPCLogPath, std::string stpQueryPCLogPath);
		void initSolver();
		void storeCacheDFS(SchedNode *node, unsigned depth, std::ofstream &output);
		void loadCacheDFS(SchedNode *node, SchedTrunk *curTrunk, unsigned &depth, std::ifstream &input, int callSite);
		bool satisfy(CacheItr *cItr, ConstraintManager *cm, std::vector<const Array*> &objects, void *addr, unsigned nbytes, const char *name);
		bool satisfyWithCache(CacheItr *cItr, ConstraintManager *cm, std::vector<const Array*> &objects, Char_Ptr_To_Int_Map &cache, void *addr, unsigned nbytes, const char *name);
		bool solve(ConstraintManager *cm);
		void printCst(ConstraintManager *cm);
		void dumpMemoizedSched(int winNum);
		void printSchedTrunk(SchedTrunk *tr);
		bool lookupConcreteCache(ConstraintManager *cm, std::vector<const Array*> &objects, std::vector< std::vector<unsigned char> > &values, Char_Ptr_To_Int_Map &cache);

                // YJF: copied from above functions; not sure why these names
		int countSatisfiedWithCache(CacheItr *cItr, ConstraintManager *cm, std::vector<const Array*> &objects, Char_Ptr_To_Int_Map &cache, void *addr, unsigned nbytes, const char *name);
		int countSatisfiedWithConcrete(ConstraintManager *cm, std::vector<const Array*> &objects, std::vector< std::vector<unsigned char> > &values, Char_Ptr_To_Int_Map &cache);


	public:
		SchedCache();
		SchedTrunk *queryNextTrunk(CacheItr *cItr, void *addr, unsigned nbytes, const char *name, int replayerTid);
		int queryOld(void *addr, unsigned nbytes, const char *name);
		void dumpMemoizedSchedToDisk(int schedId, int winNum);
		unsigned size();
		bool simplifyExpr(MemoizedSched *ms);
		void mergeSched(MemoizedSched *ms);
		void loadCache();
		void storeCache(int winNum);
		void optimizeCache();
		SchedNode *getRoot();
		bool constraintsDistinctWithObjects(ConstraintManager *cm1, std::vector<const Array*> &objects1, 
	ConstraintManager *cm2, std::vector<const Array*> &objects2);
		bool constraintsDistinct(SchedTrunk *tr1, SchedTrunk *tr2);
		bool schedTrunkDistinct(SchedTrunk *tr1, SchedTrunk *tr2);
		void createMemoizedSched();
		void fillSymObjs(KeyType &key, std::vector<const Array*> &objects);
		
		// The two APIs for recording mode. Schedule layout is transparent to these two functions, so the KLEE engine does not need to know
		// schedule layout in SchedCache.
		// I will leave the implementation with schedule trunks (decision tree) because implementing these two functions with my current schedule
		// implementation will cost some time, and all the other interface are tested, so why not leave the implementation to the decision tree.
		int addTid(char *tidInfo);
		int addConstraint(unsigned trackId, ref<Expr> expr);
		void addSymVar(const Array *arr);
		void adjustPriority(SchedTrunk *curTrunk, unsigned curTrackIndex, int brokenReason);
		void addSchedTrack(SchedTrunk *curTrunk);
		SchedTrack *getTrackByIndex(unsigned index);
		void delRPTidFromAllTracks(int replayTid);
	};
}
#endif

